# coding: utf-8
from __future__ import division
from utils.utils import *
import os
import sys
import time
import datetime
import argparse
import tqdm

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
import torch.optim as optim
from dataset.dataset import MyDataset
from models.yolov3 import yolov3
from models.yololayer import YOLOLayer


if __name__ == "__main__":
    USE_CUDA = True
    GPU_ID = [0]
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(i) for i in GPU_ID])
    device = torch.device("cuda:0" if torch.cuda.is_available() and USE_CUDA else "cpu")

    MEAN = 127.5
    IMAGE_DIM = 256
    iou_thres = 0.5
    conf_thresh = 0.85
    nms_thresh = 0.35

    cp_path = "C:\\Users\\ZiGAA\\Documents\\model\\yolov3_iter_1700000.pkl"
    dataset_root = "E:\\zhigan\\detection\\data"
    dataset_file = "E:\\zhigan\\detection\\labels.txt"

    # cp_path = r"C:\Users\ZiGAA\Documents\model\yolov3_gui_400000.pkl"
    # dataset_root = r"E:\BaiduNetdiskDownload\duojinggui\train"
    # dataset_file = r"E:\BaiduNetdiskDownload\duojinggui\train\files.txt"

    num_classes = 1
    class_names = ["A", "B", "C"]

    # ANCHORS
    with open("config/coco_anchors.txt") as f:
        lines = f.readlines()
        anchors = [[int(i.split(',')[0]), int(i.split(',')[1])] for i in lines]
        anchors = [[float(w) / 416 * 256, float(h) / 416 * 256] for w, h in anchors]

    net = yolov3(num_classes)
    head = [YOLOLayer(anchors[0:3], num_classes),
            YOLOLayer(anchors[3:6], num_classes),
            YOLOLayer(anchors[6:9], num_classes)]

    net = load_weights(cp_path, net)

    net.to(device)
    net.eval()

    head = [i.to(device) for i in head]
    head = [i.eval() for i in head]

    print("Compute mAP...")
    # Get dataloader
    dataset = MyDataset(dataset_root, dataset_file)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=1, shuffle=False, num_workers=1, collate_fn=dataset.collate_fn)

    labels = []
    sample_metrics = []  # List of tuples (TP, confs, pred)
    for batch_i, (imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc="Detecting objects")):
        # Extract labels
        labels += targets[:, 1].tolist()
        # Rescale target
        targets[:, 2:] = xywh2xyxy(targets[:, 2:])
        targets[:, 2:] *= IMAGE_DIM

        imgs = imgs.to(device)

        with torch.no_grad():
            y1, y2, y3 = net(imgs)
            o1, _ = head[2](y1, img_dim=256)
            o2, _ = head[1](y2, img_dim=256)
            o3, _ = head[0](y3, img_dim=256)
            output = torch.cat((o1, o2, o3), 1)
            output = output.cpu()
            targets = targets.cpu()
            outputs = non_max_suppression(output, conf_thres=conf_thresh, nms_thres=nms_thresh)

        sample_metrics += get_batch_statistics(outputs, targets, iou_threshold=iou_thres)

    # Concatenate sample statistics
    true_positives, pred_scores, pred_labels = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
    precision, recall, AP, f1, ap_class = ap_per_class(true_positives, pred_scores, pred_labels, labels)
    print("Precison:")
    print(precision)
    print("Recall:")
    print(recall)
    print("AP:")
    print(AP)
    print("f1:")
    print(f1)

    print("Average Precisions:")
    for i, c in enumerate(ap_class):
        print(f"+ Class '{c}' ({class_names[c]}) - AP: {AP[i]}")

    print(f"mAP: {AP.mean()}")
